package com.mm.cloth.rmbg;

import ai.onnxruntime.*;
import android.content.res.AssetManager;
import android.graphics.Bitmap;
import android.graphics.Color;
import android.util.Log;

import java.io.InputStream;
import java.nio.FloatBuffer;
import java.util.Collections;

/**
 * 基于ONNX Runtime的MODNet推理助手
 * 参考提供的ONNXInferenceHelper代码逻辑
 */
public class ONNXInferenceHelper {
    private static final String TAG = "ONNXInferenceHelper";
    
    private OrtSession session;
    private OrtEnvironment env;
    
    // MODNet模型参数
    private static final int MODEL_INPUT_SIZE = 320; // MODNet标准输入尺寸
//    private static final String INPUT_NAME = "input";
    private static final String INPUT_NAME = "input.1";
    private static final String OUTPUT_NAME = "output";
    
    // ImageNet标准化参数
    private static final float[] MEAN = {0.485f, 0.456f, 0.406f};
    private static final float[] STD = {0.229f, 0.224f, 0.225f};
    
    /**
     * 初始化ONNX Runtime环境
     * @param assets AssetManager
     * @param modelPath 模型文件路径
     */
    public void initModel(AssetManager assets, String modelPath) throws Exception {
        Log.d(TAG, "初始化ONNX Runtime环境，模型路径: " + modelPath);
        
        try {
            // 创建ONNX Runtime环境
            env = OrtEnvironment.getEnvironment();
            
            // 配置会话选项
            OrtSession.SessionOptions options = new OrtSession.SessionOptions();
            options.setOptimizationLevel(OrtSession.SessionOptions.OptLevel.ALL_OPT);
            options.setIntraOpNumThreads(4); // 设置线程数
            
            // 从assets加载模型
            InputStream modelStream = assets.open(modelPath);
            byte[] modelBytes = new byte[modelStream.available()];
            modelStream.read(modelBytes);
            modelStream.close();
            
            // 创建会话
            session = env.createSession(modelBytes, options);
            
            Log.d(TAG, "ONNX Runtime模型加载成功");
            
            // 打印模型信息
            printModelInfo();
            
        } catch (Exception e) {
            Log.e(TAG, "ONNX Runtime模型初始化失败", e);
            throw e;
        }
    }
    
    /**
     * 打印模型信息
     */
    private void printModelInfo() {
        try {
            Log.d(TAG, "=== 模型信息 ===");
            Log.d(TAG, "输入节点数量: " + session.getInputNames().size());
            Log.d(TAG, "输出节点数量: " + session.getOutputNames().size());
            
            for (String inputName : session.getInputNames()) {
                Log.d(TAG, "输入节点: " + inputName);
            }
            
            for (String outputName : session.getOutputNames()) {
                Log.d(TAG, "输出节点: " + outputName);
            }
            
        } catch (Exception e) {
            Log.w(TAG, "获取模型信息失败", e);
        }
    }
    
    /**
     * 执行抠图推理
     * @param input 输入图像
     * @return 推理结果遮罩数据
     */
    public float[] predict(Bitmap input) throws Exception {
        if (session == null || env == null) {
            throw new IllegalStateException("模型未初始化");
        }
        
        Log.d(TAG, "开始MODNet推理，输入尺寸: " + input.getWidth() + "x" + input.getHeight());
        
        long startTime = System.currentTimeMillis();
        
        try {
            // 1. 预处理：调整尺寸+归一化
            float[] inputData = preprocessImage(input);
            long[] shape = {1, 3, MODEL_INPUT_SIZE, MODEL_INPUT_SIZE}; // NCHW格式
//            long[] shape = {1, 3, 768, 768}; // NCHW格式

            Log.d(TAG, "预处理完成，输入数据长度: " + inputData.length + ", 形状: [" + 
                  shape[0] + "," + shape[1] + "," + shape[2] + "," + shape[3] + "]");
            
            // 2. 构建输入张量
            OnnxTensor tensor = OnnxTensor.createTensor(env, FloatBuffer.wrap(inputData), shape);
            
            // 3. 执行推理
            long inferenceStart = System.currentTimeMillis();
            try (OrtSession.Result result = session.run(Collections.singletonMap(INPUT_NAME, tensor))) {
                long inferenceEnd = System.currentTimeMillis();

                // 4. 获取输出结果 - MODNet输出4维张量 [batch, channel, height, width]
                Object outputObject = result.get(0).getValue();
                float[] outputData = convertOutputToFloatArray(outputObject);

                long totalTime = System.currentTimeMillis() - startTime;
                Log.d(TAG, "MODNet推理完成，推理耗时: " + (inferenceEnd - inferenceStart) + "ms, 总耗时: " + totalTime + "ms");
                Log.d(TAG, "输出数据长度: " + outputData.length);

                // 5. 统计输出数据
                logOutputStatistics(outputData);

                return outputData;
                
            } finally {
                tensor.close();
            }
            
        } catch (Exception e) {
            Log.e(TAG, "MODNet推理失败", e);
            throw e;
        }
    }

    /**
     * 转换ONNX输出为float数组
     * MODNet输出格式: float[batch][channel][height][width]
     */
    private float[] convertOutputToFloatArray(Object outputObject) {
        try {
            if (outputObject instanceof float[][][][]) {
                // 4维张量格式 [batch, channel, height, width]
                float[][][][] output4D = (float[][][][]) outputObject;

                int batch = output4D.length;
                int channels = output4D[0].length;
                int height = output4D[0][0].length;
                int width = output4D[0][0][0].length;

                Log.d(TAG, "输出张量形状: [" + batch + ", " + channels + ", " + height + ", " + width + "]");

                // 转换为1维数组 (取第一个batch，第一个channel)
                float[] result = new float[height * width];
                for (int h = 0; h < height; h++) {
                    for (int w = 0; w < width; w++) {
                        result[h * width + w] = output4D[0][0][h][w]; // 取第一个batch和channel
                    }
                }

                return result;

            } else if (outputObject instanceof float[]) {
                // 已经是1维数组
                Log.d(TAG, "输出已经是1维数组格式");
                return (float[]) outputObject;

            } else {
                Log.e(TAG, "未知的输出格式: " + outputObject.getClass().getName());
                throw new RuntimeException("不支持的输出格式: " + outputObject.getClass().getName());
            }

        } catch (Exception e) {
            Log.e(TAG, "输出转换失败", e);
            throw new RuntimeException("输出转换失败", e);
        }
    }

    /**
     * 预处理图像：RGBA转RGB+归一化
     * @param bitmap 输入图像
     * @return 预处理后的float数组
     */
    private float[] preprocessImage(Bitmap bitmap) {
        Log.d(TAG, "开始图像预处理");
        
        // 1. 缩放到模型输入尺寸
        Bitmap resizedBitmap = Bitmap.createScaledBitmap(bitmap, MODEL_INPUT_SIZE, MODEL_INPUT_SIZE, true);
        
        // 2. 获取像素数据
        int[] pixels = new int[MODEL_INPUT_SIZE * MODEL_INPUT_SIZE];
        resizedBitmap.getPixels(pixels, 0, MODEL_INPUT_SIZE, 0, 0, MODEL_INPUT_SIZE, MODEL_INPUT_SIZE);
        
        // 3. 转换为NCHW格式的float数组并进行ImageNet标准化
        int totalPixels = MODEL_INPUT_SIZE * MODEL_INPUT_SIZE;
        float[] inputData = new float[3 * totalPixels]; // RGB三通道
        
        // 分别处理R、G、B三个通道
        for (int c = 0; c < 3; c++) {
            for (int h = 0; h < MODEL_INPUT_SIZE; h++) {
                for (int w = 0; w < MODEL_INPUT_SIZE; w++) {
                    int pixelIndex = h * MODEL_INPUT_SIZE + w;
                    int pixel = pixels[pixelIndex];
                    
                    float value;
                    if (c == 0) { // R通道
                        value = Color.red(pixel) / 255.0f;
                    } else if (c == 1) { // G通道
                        value = Color.green(pixel) / 255.0f;
                    } else { // B通道
                        value = Color.blue(pixel) / 255.0f;
                    }
                    
                    // ImageNet标准化: (value - mean) / std
                    value = (value - MEAN[c]) / STD[c];
                    
                    // NCHW格式存储
                    int outputIndex = c * totalPixels + pixelIndex;
                    inputData[outputIndex] = value;
                }
            }
        }
        
        // 清理临时bitmap
        if (resizedBitmap != bitmap) {
            resizedBitmap.recycle();
        }
        
        Log.d(TAG, "图像预处理完成，输出数据长度: " + inputData.length);
        
        return inputData;
    }
    
    /**
     * 统计输出数据
     */
    private void logOutputStatistics(float[] outputData) {
        if (outputData == null || outputData.length == 0) {
            Log.w(TAG, "输出数据为空");
            return;
        }
        
        float min = Float.MAX_VALUE;
        float max = Float.MIN_VALUE;
        float sum = 0;
        int foregroundCount = 0;
        
        for (float value : outputData) {
            min = Math.min(min, value);
            max = Math.max(max, value);
            sum += value;
            if (value > 0.5f) foregroundCount++;
        }
        
        float avg = sum / outputData.length;
        float foregroundRatio = foregroundCount * 100.0f / outputData.length;
        
        Log.d(TAG, "输出统计: min=" + String.format("%.3f", min) + 
                  ", max=" + String.format("%.3f", max) + 
                  ", avg=" + String.format("%.3f", avg));
        Log.d(TAG, "前景像素比例: " + String.format("%.1f", foregroundRatio) + "%");
    }
    
    /**
     * 释放资源
     */
    public void release() {
        try {
            if (session != null) {
                session.close();
                session = null;
            }
            if (env != null) {
                env.close();
                env = null;
            }
            Log.d(TAG, "ONNX Runtime资源已释放");
        } catch (Exception e) {
            Log.e(TAG, "释放ONNX Runtime资源失败", e);
        }
    }
    
    /**
     * 检查是否已初始化
     */
    public boolean isInitialized() {
        return session != null && env != null;
    }
}
